import cv2
import numpy as np
import onnxruntime
import torch

class EfficientNetV2B0_128_ONNX:
    def __init__(self, onnx_path: str, device: str = 'auto'):
        """
        Initializes EfficientNetV2B0_128 feature extractor with ONNX runtime.
        
        Args:
            onnx_path: Path to ONNX model file
            device: Compute device ('auto', 'cuda' or 'cpu')
        """
        # Set device
        if device == 'auto':
            self.device = 'cuda' if torch.cuda.is_available() else 'cpu'
        else:
            self.device = device.lower()
            
        # Initialize ONNX runtime session
        providers = ['CUDAExecutionProvider', 'CPUExecutionProvider'] if self.device == 'cuda' else ['CPUExecutionProvider']
        self.session = onnxruntime.InferenceSession(onnx_path, providers=providers)
        
        # Get model metadata
        self.input_name = self.session.get_inputs()[0].name
        self.output_name = self.session.get_outputs()[0].name
        self.input_shape = self.session.get_inputs()[0].shape
        
        print(f"EfficientNetV2B0_128 ONNX model loaded successfully!")
        print(f"Input shape: {self.input_shape}")
        print(f"Using device: {self.device}")
        
        # ImageNet normalization parameters
        self.mean = np.array([0.485, 0.456, 0.406], dtype=np.float32)
        self.std = np.array([0.229, 0.224, 0.225], dtype=np.float32)

    def preprocess(self, img: np.ndarray) -> np.ndarray:
        """
        Preprocess input image for EfficientNetV2B0_128.
        
        Args:
            img: Input image in BGR format (H,W,3)
            
        Returns:
            np.ndarray: Preprocessed image (1,3,224,224)
        """
        # Convert BGR to RGB
        img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
        
        # Resize to 224x224
        img = cv2.resize(img, (224, 224))
        
        # Convert to float32 and normalize to [0,1]
        img = img.astype(np.float32) / 255.0
        
        # Normalize with ImageNet mean and std
        img = (img - self.mean) / self.std
        
        # Transpose to CHW format and add batch dimension
        img = np.transpose(img, (2, 0, 1))
        img = np.expand_dims(img, axis=0)
        
        return img

    def extract_features(self, img: np.ndarray) -> np.ndarray:
        """
        Extract 128-dimensional feature vector from input image.
        
        Args:
            img: Input image in BGR format (H,W,3)
            
        Returns:
            np.ndarray: 128-dimensional feature vector
        """
        # Preprocess image
        img = self.preprocess(img)
        
        # Run inference
        features = self.session.run(
            [self.output_name],
            {self.input_name: img}
        )[0]
        
        # L2 normalize the features
        features = features / np.linalg.norm(features)
        
        return features
